from hetong.qwen.mode_api import send_message

def extract_info(text, query):
    # 设计提示
    # prompt = (
    #     f"\n\n{text}\n\n"
    #     f"请从给定的文本中精确抽取{query}的数据，并将其封装为JSON格式的数据。JSON应仅包含一个键值对，其中键名为'{query}'，值为文本中对应的合同金额数值。如果文本中未找到合同金额的数值，请返回一个空字符串作为值。输出格式应为：{\'key\':\'合同金额\',\'value\':\'抽取的值\'}。"
    # )
    # json_str = '{"key":"%s","value":"抽取的数据"}' % query
    json_str = '{"key":"%s","value":"%s的数据范围"}' % (query,query)
#     prompt = (
#         f'''
# \n\n{text}\n\n
# 你是一个数据抽取专家，任务是准确地从给定的合同文本中识别并提取{query}的数值。
# 请从给定的文本中精确抽取{query}的数据，并将其封装为JSON格式的数据。JSON应仅包含一个键值对，其中键名为'{query}'，值为文本中对应的{query}的数据。
# 如果文本中未找到{query}的数据,不需要补充额外的数据，直接返回空字符串的json,输出格式应为：{json_str},不要额外补充
#         '''
#     )
    prompt = (
        f'''
 \n\n合同文本:{text}\n\n
您的任务是审查提供的合同文本，并从中准确地抽取第一个出现的“{query}”。将这些信息整合成一个JSON数据格式，该数据仅包含一个键值对，其中键名固定为“{query}”，值则应是您从文本中抽取出的{query}的结合。如果文本中未明确列出{query}的具体数据或描述，请仅使用{query}作为值。输出格式应当遵循以下模式：{json_str}。

请注意，如果{query}服务范围无法确定，值应直接为{query}，不附加任何额外信息。
        '''
    )
    print('prompt:', prompt)
    response = send_message(prompt, system_prompt='你是一个数据分析人员，擅长总结与抽取实用数据')
    if 'id' in response:
        # 获取模型输出
        print(response)
        return response['choices'][0]['message']['content']
    else:
        print('Request id: %s, Status code: %s, error code: %s, error message: %s' % (
            response.id, response.model,
            response.object
        ))


if __name__ == '__main__':
    extract_info('wefwefaefaewf', 'afewawef')
